// Copyright (C) joyinter
// All rights reserved.
// 
// Author : Jiang Li
// Date   : 2004-6-4

#include "aes.h"
#include <windows.h>

#define AES_MIN_KEY_SIZE	16
#define AES_MAX_KEY_SIZE	32

#define AES_BLOCK_SIZE		16

//##ModelId=44C4B5FD0203
typedef DWORD u32;
//##ModelId=44C4B5FD0213
typedef BYTE  u8;

static inline 
u32 generic_rotr32 (const u32 x, const unsigned bits)
{
	const unsigned n = bits % 32;
	return (x >> n) | (x << (32 - n));
}

static inline 
u32 generic_rotl32 (const u32 x, const unsigned bits)
{
	const unsigned n = bits % 32;
	return (x << n) | (x >> (32 - n));
}

#define rotl generic_rotl32
#define rotr generic_rotr32

/*
 * #define to_byte(x, nr) ((unsigned char)((x) >> (nr*8))) 
 */
inline static u8
to_byte(const u32 x, const unsigned n)
{
	return (u8)(x >> (n << 3));
}

//#define u32_in(x) le32_to_cpu(*(const u32 *)(x))
//#define u32_out(to, from) (*(u32 *)(to) = cpu_to_le32(from))
#define u32_in(x) (*(const u32 *)(x))
#define u32_out(to, from) (*(u32 *)(to) = (from))

//##ModelId=44C4B5FD0222
struct aes_ctx {
	//##ModelId=44C4B5FD0224
	int key_length;
	//##ModelId=44C4B5FD0233
	u32 E[60];
	//##ModelId=44C4B5FD0238
	u32 D[60];
};

#define E_KEY ctx->E
#define D_KEY ctx->D

//#define GEN_TABLE
#ifdef GEN_TABLE

static u8 pow_tab[256];
static u8 log_tab[256];
static u8 sbx_tab[256];
static u8 isb_tab[256];
static u32 rco_tab[10];
static u32 ft_tab[4][256];
static u32 it_tab[4][256];

static u32 fl_tab[4][256];
static u32 il_tab[4][256];

#else

#include "aestab.h"

#endif

static inline u8
f_mult (u8 a, u8 b)
{
	u8 aa = log_tab[a], cc = aa + log_tab[b];

	return pow_tab[cc + (cc < aa ? 1 : 0)];
}

#define ff_mult(a,b)    (a && b ? f_mult(a, b) : 0)

#define f_rn(bo, bi, n, k)					\
    bo[n] =  ft_tab[0][to_byte(bi[n],0)] ^				\
             ft_tab[1][to_byte(bi[(n + 1) & 3],1)] ^		\
             ft_tab[2][to_byte(bi[(n + 2) & 3],2)] ^		\
             ft_tab[3][to_byte(bi[(n + 3) & 3],3)] ^ *(k + n)

#define i_rn(bo, bi, n, k)					\
    bo[n] =  it_tab[0][to_byte(bi[n],0)] ^				\
             it_tab[1][to_byte(bi[(n + 3) & 3],1)] ^		\
             it_tab[2][to_byte(bi[(n + 2) & 3],2)] ^		\
             it_tab[3][to_byte(bi[(n + 1) & 3],3)] ^ *(k + n)

#define ls_box(x)				\
    ( fl_tab[0][to_byte(x, 0)] ^			\
      fl_tab[1][to_byte(x, 1)] ^			\
      fl_tab[2][to_byte(x, 2)] ^			\
      fl_tab[3][to_byte(x, 3)] )

#define f_rl(bo, bi, n, k)					\
    bo[n] =  fl_tab[0][to_byte(bi[n],0)] ^				\
             fl_tab[1][to_byte(bi[(n + 1) & 3],1)] ^		\
             fl_tab[2][to_byte(bi[(n + 2) & 3],2)] ^		\
             fl_tab[3][to_byte(bi[(n + 3) & 3],3)] ^ *(k + n)

#define i_rl(bo, bi, n, k)					\
    bo[n] =  il_tab[0][to_byte(bi[n],0)] ^				\
             il_tab[1][to_byte(bi[(n + 3) & 3],1)] ^		\
             il_tab[2][to_byte(bi[(n + 2) & 3],2)] ^		\
             il_tab[3][to_byte(bi[(n + 1) & 3],3)] ^ *(k + n)

#ifdef GEN_TABLE
#include "saesgentab.h"
#endif // GEN_TABLE

#define star_x(x) (((x) & 0x7f7f7f7f) << 1) ^ ((((x) & 0x80808080) >> 7) * 0x1b)

#define imix_col(y,x)       \
    u   = star_x(x);        \
    v   = star_x(u);        \
    w   = star_x(v);        \
    t   = w ^ (x);          \
   (y)  = u ^ v ^ w;        \
   (y) ^= rotr(u ^ t,  8) ^ \
          rotr(v ^ t, 16) ^ \
          rotr(t,24)

/* initialise the key schedule from the user supplied key */

#define loop4(i)                                    \
{   t = rotr(t,  8); t = ls_box(t) ^ rco_tab[i];    \
    t ^= E_KEY[4 * i];     E_KEY[4 * i + 4] = t;    \
    t ^= E_KEY[4 * i + 1]; E_KEY[4 * i + 5] = t;    \
    t ^= E_KEY[4 * i + 2]; E_KEY[4 * i + 6] = t;    \
    t ^= E_KEY[4 * i + 3]; E_KEY[4 * i + 7] = t;    \
}

#define loop6(i)                                    \
{   t = rotr(t,  8); t = ls_box(t) ^ rco_tab[i];    \
    t ^= E_KEY[6 * i];     E_KEY[6 * i + 6] = t;    \
    t ^= E_KEY[6 * i + 1]; E_KEY[6 * i + 7] = t;    \
    t ^= E_KEY[6 * i + 2]; E_KEY[6 * i + 8] = t;    \
    t ^= E_KEY[6 * i + 3]; E_KEY[6 * i + 9] = t;    \
    t ^= E_KEY[6 * i + 4]; E_KEY[6 * i + 10] = t;   \
    t ^= E_KEY[6 * i + 5]; E_KEY[6 * i + 11] = t;   \
}

#define loop8(i)                                    \
{   t = rotr(t,  8); ; t = ls_box(t) ^ rco_tab[i];  \
    t ^= E_KEY[8 * i];     E_KEY[8 * i + 8] = t;    \
    t ^= E_KEY[8 * i + 1]; E_KEY[8 * i + 9] = t;    \
    t ^= E_KEY[8 * i + 2]; E_KEY[8 * i + 10] = t;   \
    t ^= E_KEY[8 * i + 3]; E_KEY[8 * i + 11] = t;   \
    t  = E_KEY[8 * i + 4] ^ ls_box(t);    \
    E_KEY[8 * i + 12] = t;                \
    t ^= E_KEY[8 * i + 5]; E_KEY[8 * i + 13] = t;   \
    t ^= E_KEY[8 * i + 6]; E_KEY[8 * i + 14] = t;   \
    t ^= E_KEY[8 * i + 7]; E_KEY[8 * i + 15] = t;   \
}

static bool aes_set_key(void *ctx_arg,
    const u8 *in_key, unsigned int key_len)
{
	struct aes_ctx *ctx = (aes_ctx*)ctx_arg;
	u32 i, t, u, v, w;

	if (key_len != 16 && key_len != 24 && key_len != 32)
    {
		return false;
	}

	ctx->key_length = key_len;

	E_KEY[0] = u32_in (in_key);
	E_KEY[1] = u32_in (in_key + 4);
	E_KEY[2] = u32_in (in_key + 8);
	E_KEY[3] = u32_in (in_key + 12);

	switch (key_len)
    {
	case 16:
		t = E_KEY[3];
		for (i = 0; i < 10; ++i)
			loop4 (i);
		break;

	case 24:
		E_KEY[4] = u32_in (in_key + 16);
		t = E_KEY[5] = u32_in (in_key + 20);
		for (i = 0; i < 8; ++i)
			loop6 (i);
		break;

	case 32:
		E_KEY[4] = u32_in (in_key + 16);
		E_KEY[5] = u32_in (in_key + 20);
		E_KEY[6] = u32_in (in_key + 24);
		t = E_KEY[7] = u32_in (in_key + 28);
		for (i = 0; i < 7; ++i)
			loop8 (i);
		break;
	}

	D_KEY[0] = E_KEY[0];
	D_KEY[1] = E_KEY[1];
	D_KEY[2] = E_KEY[2];
	D_KEY[3] = E_KEY[3];

	for (i = 4; i < key_len + 24; ++i) {
		imix_col (D_KEY[i], E_KEY[i]);
	}

	return true;
}

/* encrypt a block of text */

#define f_nround(bo, bi, k) \
    f_rn(bo, bi, 0, k);     \
    f_rn(bo, bi, 1, k);     \
    f_rn(bo, bi, 2, k);     \
    f_rn(bo, bi, 3, k);     \
    k += 4

#define f_lround(bo, bi, k) \
    f_rl(bo, bi, 0, k);     \
    f_rl(bo, bi, 1, k);     \
    f_rl(bo, bi, 2, k);     \
    f_rl(bo, bi, 3, k)

static void aes_encrypt(void *ctx_arg, u8 *out, const u8 *in)
{
	const struct aes_ctx *ctx = (aes_ctx*)ctx_arg;
	u32 b0[4], b1[4];
	const u32 *kp = E_KEY + 4;

	b0[0] = u32_in (in) ^ E_KEY[0];
	b0[1] = u32_in (in + 4) ^ E_KEY[1];
	b0[2] = u32_in (in + 8) ^ E_KEY[2];
	b0[3] = u32_in (in + 12) ^ E_KEY[3];

    const u32 *kpend = kp + ctx->key_length + 16;

    do
    {
		f_nround (b1, b0, kp);
		f_nround (b0, b1, kp);
    }
    while(kp < kpend);
/*
	if (ctx->key_length > 24) {
		f_nround (b1, b0, kp);
		f_nround (b0, b1, kp);
	}

	if (ctx->key_length > 16) {
		f_nround (b1, b0, kp);
		f_nround (b0, b1, kp);
	}

	f_nround (b1, b0, kp);
	f_nround (b0, b1, kp);
	f_nround (b1, b0, kp);
	f_nround (b0, b1, kp);
	f_nround (b1, b0, kp);
	f_nround (b0, b1, kp);
	f_nround (b1, b0, kp);
	f_nround (b0, b1, kp);
*/
	f_nround (b1, b0, kp);
	f_lround (b0, b1, kp);
	u32_out (out, b0[0]);
	u32_out (out + 4, b0[1]);
	u32_out (out + 8, b0[2]);
	u32_out (out + 12, b0[3]);
}

/* decrypt a block of text */

#define i_nround(bo, bi, k) \
    i_rn(bo, bi, 0, k);     \
    i_rn(bo, bi, 1, k);     \
    i_rn(bo, bi, 2, k);     \
    i_rn(bo, bi, 3, k);     \
    k -= 4

#define i_lround(bo, bi, k) \
    i_rl(bo, bi, 0, k);     \
    i_rl(bo, bi, 1, k);     \
    i_rl(bo, bi, 2, k);     \
    i_rl(bo, bi, 3, k)

static void aes_decrypt(void *ctx_arg, u8 *out, const u8 *in)
{
	const struct aes_ctx *ctx = (aes_ctx*)ctx_arg;
	u32 b0[4], b1[4];
	const int key_len = ctx->key_length;
	const u32 *kp = D_KEY + key_len + 20;

	b0[0] = u32_in (in) ^ E_KEY[key_len + 24];
	b0[1] = u32_in (in + 4) ^ E_KEY[key_len + 25];
	b0[2] = u32_in (in + 8) ^ E_KEY[key_len + 26];
	b0[3] = u32_in (in + 12) ^ E_KEY[key_len + 27];

    const u32 *kpend = kp - key_len - 16;

    do
    {
		i_nround (b1, b0, kp);
		i_nround (b0, b1, kp);
    }
    while(kp > kpend);
/*
    if (key_len > 24) {
		i_nround (b1, b0, kp);
		i_nround (b0, b1, kp);
	}

	if (key_len > 16) {
		i_nround (b1, b0, kp);
		i_nround (b0, b1, kp);
	}

	i_nround (b1, b0, kp);
	i_nround (b0, b1, kp);
	i_nround (b1, b0, kp);
	i_nround (b0, b1, kp);
	i_nround (b1, b0, kp);
	i_nround (b0, b1, kp);
	i_nround (b1, b0, kp);
	i_nround (b0, b1, kp);
*/
	i_nround (b1, b0, kp);
	i_lround (b0, b1, kp);

	u32_out (out, b0[0]);
	u32_out (out + 4, b0[1]);
	u32_out (out + 8, b0[2]);
	u32_out (out + 12, b0[3]);
}


//##ModelId=44C4B5FD01C5
CJoyAes::CJoyAes(void* pKey, UINT nSize)
{
    if(pKey == NULL)
    {
        m_Context.key_length = 0;
    }
    else
    {
        SetKey(pKey, nSize);
    }
}

//##ModelId=44C4B5FD01CA
const UINT* CJoyAes::GetLegalKeySizes()
{
    static UINT keySizes[] = { 32, 24, 16, 0 };
    return keySizes;
}

//##ModelId=44C4B5FD01CC
UINT CJoyAes::GetBlockSize()
{
    return AES_BLOCK_SIZE;
}

//##ModelId=44C4B5FD01CE
UINT CJoyAes::GetKeySize()
{
    return m_Context.key_length;
}

//##ModelId=44C4B5FD01D5
bool CJoyAes::SetKey(void* pKey, UINT nSize)
{
    return aes_set_key(&m_Context, (const u8*)pKey, nSize);
}

//##ModelId=44C4B5FD01D9
bool CJoyAes::Encrypt(PVOID pDst, LPCVOID pSrc, size_t nLength)
{
    if(nLength%AES_BLOCK_SIZE != 0)
    {
        return false;
    }
    void* pEnd  = (BYTE*)pSrc + nLength;
    for( ; pSrc < pEnd; pSrc = (BYTE*)pSrc+AES_BLOCK_SIZE, pDst = (BYTE*)pDst+AES_BLOCK_SIZE)
    {
        aes_encrypt(&m_Context, (u8*)pDst, (u8*)pSrc);
    }
    return true;
}

//##ModelId=44C4B5FD01E5
bool CJoyAes::Decrypt(PVOID pDst, LPCVOID pSrc, size_t nLength)
{
    if(nLength%AES_BLOCK_SIZE != 0)
    {
        return false;
    }
    void* pEnd  = (BYTE*)pSrc + nLength;
    for( ; pSrc < pEnd; pSrc = (BYTE*)pSrc+AES_BLOCK_SIZE, pDst = (BYTE*)pDst+AES_BLOCK_SIZE)
    {
        aes_decrypt(&m_Context, (u8*)pDst, (u8*)pSrc);
    }
    return true;
}
